# -*- coding: utf-8 -*-
# ===========================================
# @Time    : 2021/8/24 上午11:40
# @Author  : shutao
# @FileName: resnet_base.py
# @remark  : 
# 
# @Software: PyCharm
# Github 　： https://github.com/NameLacker
# ===========================================

import os

import paddle
import paddle.nn as nn
from modules.utils import LRScheduler

from modules.exp.base_exp import BaseExp


class Exp(BaseExp):
    """The experiment of resnet."""

    def __init__(self):
        super().__init__()

        # -------------------- model config --------------------
        self.num_classes = 10

        # -------------------- dataloader config --------------------
        self.data_num_workers = 0

        # -------------------- transform config --------------------
        self.enable_mixup = True

        # --------------------  training config --------------------
        self.max_iters = 60000
        self.max_epoch = 30
        self.scheduler = "PolynomialDecay"
        self.min_lr_ratio = 0.05
        self.base_lr = 0.01
        self.power = 0.9

        self.weight_decay = 5e-4
        self.momentum = 0.9
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

        # --------------------  testing config --------------------
        self.test_size = (1, 3, 224, 224)
        self.test_conf = 0.01
        self.nmsthre = 0.65

    def get_model(self):
        """
        构造网络模型
        :return:
        """
        from paddle.vision import resnet18

        def init_resnet(m):
            if isinstance(m, nn.BatchNorm2D):
                m.__setattr__("_momentum", 0.03)
                m.__setattr__("_epsilon", 1e-3)

        if getattr(self, "model", None) is None:
            self.model = resnet18(pretrained=False, num_classes=self.num_classes)

        self.model.apply(init_resnet)
        return self.model

    def get_data_loader(self, batch_size, no_aug=False):
        """
        构造数据读取器
        :param batch_size:
        :param no_aug:
        :return:
        """
        from modules.data.datasets import MyDataset
        td = MyDataset()
        td_loader = paddle.io.DataLoader(td, batch_size=batch_size, shuffle=True, drop_last=True)
        return td_loader, len(td_loader)

    def get_loss_function(self, preds, loss, **kwargs) -> paddle.Tensor:
        pass

    def get_optimizer(self, scheduler, parameters):
        opt = paddle.optimizer.Momentum(learning_rate=scheduler,
                                        parameters=parameters, weight_decay=self.weight_decay, momentum=self.momentum)
        return opt

    def get_lr_scheduler(self, lr, iters_per_epoch, **kwargs):
        lr_scheduler = LRScheduler(self.scheduler,
                                   self.base_lr,
                                   max_iters=self.max_iters,
                                   power=self.power
                                   )
        scheduler = lr_scheduler.lr_func
        return scheduler

    def get_training_subject(self, data, loss_func):
        pass

    def get_evaluator(self, batch_size):
        return None

    def eval(self, model, evaluator, weights):
        pass
